<template>
  <q-table v-bind="$attrs" :rows="processedRows">
    <!-- 透传所有插槽 -->
    <template v-for="(_, slot) in $slots" #[slot]="scope">
      <slot :name="slot" v-bind="scope" />
    </template>

    <!-- 表头选择框 -->
    <template #header="props">
      <q-tr :props="props">
        <q-th v-if="$attrs.selection" auto-width>
          <q-checkbox v-model="isAllSelected" @update:model-value="handleSelectAll" />
        </q-th>
        <q-th v-for="col in props.cols" :key="col.name" :props="props">
          {{ col.label }}
        </q-th>
      </q-tr>
    </template>

    <!-- 自定义 body 实现 -->
    <template #body="props">
      <q-tr :props="props">
        <!-- 选择框列 -->
        <q-td v-if="$attrs.selection" auto-width>
          <q-checkbox
            v-model="selectedMap[props.row[rowKey]]"
            @update:model-value="(val) => handleSelect(props.row, val)"
          />
        </q-td>

        <!-- 数据列 -->
        <q-td v-for="(col, colIndex) in props.cols" :key="col.name" :props="props">
          <!-- 第一列特殊处理 -->
          <template v-if="isFirstColumn(colIndex)">
            <div :style="getIndentStyle(props.row)" class="row items-center">
              <!-- 展开按钮或占位符 -->
              <div
                class="expand-button-container"
                style="width: 24px; display: flex; justify-content: center"
              >
                <q-icon
                  v-if="props.row.__hasChildren"
                  :name="props.row.__expanded ? 'expand_more' : 'chevron_right'"
                  class="cursor-pointer"
                  size="1.2em"
                  @click.stop="toggleRow(props.row)"
                />
              </div>

              <!-- 列内容 -->
              <slot :name="`body-cell-${col.name}`" v-bind="{ ...props, col }">
                {{ formatColumnValue(col, props.row[col.field]) }}
              </slot>
            </div>
          </template>

          <!-- 其他列 -->
          <template v-else>
            <slot :name="`body-cell-${col.name}`" v-bind="{ ...props, col }">
              {{ formatColumnValue(col, props.row[col.field]) }}
            </slot>
          </template>
        </q-td>
      </q-tr>
    </template>
  </q-table>
</template>

<script>
import { ref, watch, computed } from 'vue'

export default {
  name: 'TreeTable',

  props: {
    rows: {
      type: Array,
      required: true,
    },
    rowKey: {
      type: String,
      default: 'id',
    },
    childrenKey: {
      type: String,
      default: 'children',
    },
    indentSize: {
      type: Number,
      default: 24,
    },
    defaultExpandAll: {
      type: Boolean,
      default: false,
    },
    expandedKeys: {
      type: Array,
      default: () => [],
    },
  },

  emits: ['update:expandedKeys', 'update:selected'],

  setup(props, { emit }) {
    const processedRows = ref([])
    const selectedMap = ref({})
    const expandedKeys = ref(new Set(props.expandedKeys))

    // 判断是否为第一列（排除选择列）
    const isFirstColumn = (index) => {
      const hasSelection = !!props.$attrs?.selection
      return hasSelection ? index === 0 : index === 0
    }

    // 获取缩进样式
    const getIndentStyle = (row) => ({
      'padding-left': `${row.__level * props.indentSize}px`,
      position: 'relative',
    })

    // 处理树形数据
    const flattenData = (nodes, level = 0) => {
      return nodes.reduce((acc, node) => {
        const hasChildren = node[props.childrenKey]?.length > 0
        const expanded = expandedKeys.value.has(node[props.rowKey]) || props.defaultExpandAll

        const newNode = {
          ...node,
          __level: level,
          __hasChildren: hasChildren,
          __expanded: expanded,
        }

        acc.push(newNode)

        if (expanded && hasChildren) {
          acc.push(...flattenData(node[props.childrenKey], level + 1))
        }

        return acc
      }, [])
    }

    // 切换展开状态
    const toggleRow = (row) => {
      row.__expanded = !row.__expanded
      if (row.__expanded) {
        expandedKeys.value.add(row[props.rowKey])
      } else {
        expandedKeys.value.delete(row[props.rowKey])
      }
      emit('update:expandedKeys', Array.from(expandedKeys.value))
      processedRows.value = flattenData(props.rows)
    }

    // 获取所有行的 key
    const getAllKeys = (nodes) => {
      return nodes.reduce((keys, node) => {
        keys.push(node[props.rowKey])
        if (node[props.childrenKey]?.length) {
          keys.push(...getAllKeys(node[props.childrenKey]))
        }
        return keys
      }, [])
    }

    // 获取节点的所有子节点 key
    const getChildrenKeys = (node) => {
      const keys = []
      const traverse = (n) => {
        if (n[props.childrenKey]?.length) {
          n[props.childrenKey].forEach((child) => {
            keys.push(child[props.rowKey])
            traverse(child)
          })
        }
      }
      traverse(node)
      return keys
    }

    // 获取节点的所有父节点
    const getParentNode = (key, nodes, parent = null) => {
      for (const node of nodes) {
        if (node[props.rowKey] === key) return parent
        if (node[props.childrenKey]?.length) {
          const found = getParentNode(key, node[props.childrenKey], node)
          if (found) return found
        }
      }
      return null
    }

    // 检查节点的所有子节点是否都被选中
    const areAllChildrenSelected = (node) => {
      const childrenKeys = getChildrenKeys(node)
      return childrenKeys.length > 0 && childrenKeys.every((key) => selectedMap.value[key])
    }

    // 处理单个选择
    const handleSelect = (row, val) => {
      // 更新当前节点
      selectedMap.value[row[props.rowKey]] = val

      // 更新子节点
      const childrenKeys = getChildrenKeys(row)
      childrenKeys.forEach((key) => {
        selectedMap.value[key] = val
      })

      // 更新父节点
      let parent = getParentNode(row[props.rowKey], props.rows)
      while (parent) {
        selectedMap.value[parent[props.rowKey]] = areAllChildrenSelected(parent)
        parent = getParentNode(parent[props.rowKey], props.rows)
      }

      // 发出选中值更新事件
      emitSelectedUpdate()
    }

    // 处理全选
    const handleSelectAll = (val) => {
      const allKeys = getAllKeys(props.rows)
      allKeys.forEach((key) => {
        selectedMap.value[key] = val
      })
      emitSelectedUpdate()
    }

    // 发出选中值更新事件
    const emitSelectedUpdate = () => {
      const selectedRows = processedRows.value.filter((row) => selectedMap.value[row[props.rowKey]])
      emit('update:selected', selectedRows)
    }

    // 计算是否全选
    const isAllSelected = computed(() => {
      const allKeys = getAllKeys(props.rows)
      return allKeys.length > 0 && allKeys.every((key) => selectedMap.value[key])
    })

    // 监听数据变化
    watch(
      () => props.rows,
      (newVal) => {
        processedRows.value = flattenData(newVal)
      },
      { immediate: true, deep: true },
    )

    return {
      processedRows,
      selectedMap,
      isAllSelected,
      handleSelect,
      handleSelectAll,
      toggleRow,
      isFirstColumn,
      getIndentStyle,
      formatColumnValue: (col, value) => {
        if (col.format && typeof col.format === 'function') {
          return col.format(value)
        }
        return value
      },
    }
  },
}
</script>

<style scoped>
/* 展开图标动画 */
.q-icon {
  transition: transform 0.2s ease;
}

/* 层级连接线 */
.q-td:first-child div:before {
  content: '';
}

/* 最后层级的连线隐藏 */
.q-tr:last-child .q-td:first-child div:before {
  display: none;
}
</style>
